import { ExtractChunkData } from '@cherrystudio/embedjs-interfaces'
import axiosProxy from '@main/services/AxiosProxy'
import { KnowledgeBaseParams } from '@types'

import BaseReranker from './BaseReranker'

/**
 * 基于 VoyageAI Rerank API 的重排序器
 * 继承自 BaseReranker，实现具体的 rerank 方法
 */
export default class VoyageReranker extends BaseReranker {
  /**
   * 构造函数，初始化重排序器
   * @param base 知识库参数
   */
  constructor(base: KnowledgeBaseParams) {
    super(base)
  }

  /**
   * 重排序方法
   * @param query 查询字符串
   * @param searchResults 检索到的文本块数组
   * @returns 排序后的文本块数组
   */
  public rerank = async (query: string, searchResults: ExtractChunkData[]): Promise<ExtractChunkData[]> => {
    // 获取重排序服务的 URL
    const url = this.getRerankUrl()

    // 构建请求体
    const requestBody = {
      model: this.base.rerankModel, // 使用的重排序模型
      query,
      documents: searchResults.map((doc) => doc.pageContent), // 提取所有文本内容
      top_k: this.base.topN, // 返回前K条
      return_documents: false, // 是否返回文档内容
      truncation: true // 是否启用截断
    }

    try {
      // 调用 VoyageAI Rerank API，发送 POST 请求
      const { data } = await axiosProxy.axios.post(url, requestBody, {
        headers: {
          ...this.defaultHeaders()
        }
      })

      // 获取 API 返回的重排序结果
      const rerankResults = data.data
      // 根据分数对原始检索结果排序
      return this.getRerankResult(searchResults, rerankResults)
    } catch (error: any) {
      // 格式化错误信息并输出
      const errorDetails = this.formatErrorMessage(url, error, requestBody)

      console.error('Voyage Reranker API Error:', errorDetails)
      throw new Error(`重排序请求失败: ${error.message}\n请求详情: ${errorDetails}`)
    }
  }
}
